import casadi as ca
import numpy as np
import matplotlib.pyplot as plt
import time

# opti = ca.Opti()

# x = 0
# # dx = 0
# # u = 0
# T = 0.1
# t = ca.linspace(0, T, 20)

# a0 = opti.variable()
# a1 = opti.variable()
# a2 = opti.variable()
# a3 = opti.variable()

# x = a0 + a1*t + a2*(t**2) + a3*(t**3)
# dx = a1 + a2*t + a3*(t**2)
# u = a2 + a3*t
# x0 = np.pi/3
# x1 = -np.pi/3
# dx0 = 1
# dx1 = 0

# opti.subject_to(a0 == x0)
# opti.subject_to(a1 == dx0)
# opti.subject_to(a2 == -(T**(-2))*((3*(x0 - x1))+(T*(2*dx0 + dx1))))
# opti.subject_to(a3 == (T**(-3))*((2*(x0 - x1))+(T*(dx0 + dx1))))

# # opti.subject_to(ca.jacobian(dx, t) == u)
# # opti.subject_to(ca.jacobian(x, t) == dx)

# # opti.subject_to(x[-1] == 0)
# # opti.subject_to(x[0] == 5)
# # opti.subject_to(dx[-1] == 0)
# # opti.subject_to(dx[0] == 0)
# # opti.subject_to(u[-1] == 0)
# # opti.subject_to(u[0] == 0)


# p_opts = {"expand":True}
# s_opts = {"max_iter": 500}
# opti.solver("ipopt",p_opts, s_opts)

# # try:
# fig = plt.figure()
# # debug = fig.add_subplot(121, autoscale_on=False, xlim=(0, 2), ylim=(-50, 50))
# # debug.grid()

# ax = fig.add_subplot(111, autoscale_on=False, xlim=(0, T), ylim=(-5, 5))
# ax.grid()

# # def callback(x):
# #     return debug.plot(t, opti.debug.value(x))

# # opti.callback(lambda i: callback(opti.debug.value(x)))
# try:
#     sol = opti.solve() 
#     print(sol.stats()["iter_count"])

#     # print(sol.value(x))
#     # t = np.linspace(0, 2, 20)
#     # f = - 1 - 2*t - 3*(t**2) - 2*(t**3) - 5*(t**4) - 6*(t**5)

#     ax.plot(t, sol.value(x), 'r-' , label='x')
#     ax.plot(t, sol.value(dx), 'b-' , label='dx')
#     ax.plot(t, sol.value(u), 'g-' , label='u')

#     plt.show()
# except:
#     debug_x = opti.debug.value(x)
#     debug_dx = opti.debug.value(dx)
#     debug_u = opti.debug.value(u)

#     print(debug_x)
#     print(debug_dx)
#     print(debug_u)

#     ax.plot(t, debug_x, 'r-' , label='x')
#     ax.plot(t, debug_dx, 'b-' , label='dx')
#     ax.plot(t, debug_u, 'g-' , label='u')

#     plt.show()





# Degree of interpolating polynomial
d = 3

# Get collocation points
tau_root = np.append(0, ca.collocation_points(d, 'legendre'))

# Coefficients of the collocation equation
C = np.zeros((d+1,d+1))

# Coefficients of the continuity equation
D = np.zeros(d+1)

# Coefficients of the quadrature function
B = np.zeros(d+1)

# Construct polynomial basis
for j in range(d+1):
    # Construct Lagrange polynomials to get the polynomial basis at the collocation point
    p = np.poly1d([1])
    for r in range(d+1):
        if r != j:
            p *= np.poly1d([1, -tau_root[r]]) / (tau_root[j]-tau_root[r])

    # Evaluate the polynomial at the final time to get the coefficients of the continuity equation
    D[j] = p(1.0)

    # Evaluate the time derivative of the polynomial at all collocation points to get the coefficients of the continuity equation
    pder = np.polyder(p)
    for r in range(d+1):
        C[j,r] = pder(tau_root[r])

    # Evaluate the integral of the polynomial to get the coefficients of the quadrature function
    pint = np.polyint(p)
    B[j] = pint(1.0)

print(tau_root)
print(C)
print(D)
print(B)

# # Time horizon
# T = 10.

# # Declare model variables
# x1 = ca.SX.sym('x1')
# x2 = ca.SX.sym('x2')
# x = ca.vertcat(x1, x2)
# u = ca.SX.sym('u')

# # Model equations
# xdot = ca.vertcat((1-x2**2)*x1 - x2 + u, x1)

# # Objective term
# L = x1**2 + x2**2 + u**2

# # Continuous time dynamics
# f = ca.Function('f', [x, u], [xdot, L], ['x', 'u'], ['xdot', 'L'])

# # Control discretization
# N = 5 # number of control intervals
# h = T/N

# # Start with an empty NLP
# w=[]
# w0 = []
# lbw = []
# ubw = []
# J = 0
# g=[]
# lbg = []
# ubg = []

# # For plotting x and u given w
# x_plot = []
# u_plot = []

# # "Lift" initial conditions
# Xk = ca.MX.sym('X0', 2)
# w.append(Xk)
# lbw.append([0, 1])
# ubw.append([0, 1])
# w0.append([0, 1])
# x_plot.append(Xk)

# # Formulate the NLP
# for k in range(N):
#     # New NLP variable for the control
#     Uk = ca.MX.sym('U_' + str(k))
#     w.append(Uk)
#     lbw.append([-1])
#     ubw.append([1])
#     w0.append([0])
#     u_plot.append(Uk)

#     # State at collocation points
#     Xc = []
#     for j in range(d):
#         Xkj = ca.MX.sym('X_'+str(k)+'_'+str(j), 2)
#         Xc.append(Xkj)
#         w.append(Xkj)
#         lbw.append([-0.25, -np.inf])
#         ubw.append([np.inf,  np.inf])
#         w0.append([0, 0])

#     # Loop over collocation points
#     Xk_end = D[0]*Xk
#     for j in range(1,d+1):
#        # Expression for the state derivative at the collocation point
#        xp = C[0,j]*Xk
#        for r in range(d): xp = xp + C[r+1,j]*Xc[r]

#        # Append collocation equations
#        fj, qj = f(Xc[j-1],Uk)
#        g.append(h*fj - xp)
#        lbg.append([0, 0])
#        ubg.append([0, 0])

#        # Add contribution to the end state
#        Xk_end = Xk_end + D[j]*Xc[j-1];

#        # Add contribution to quadrature function
#        J = J + B[j]*qj*h

#     # New NLP variable for state at end of interval
#     Xk = ca.MX.sym('X_' + str(k+1), 2)
#     w.append(Xk)
#     lbw.append([-0.25, -np.inf])
#     ubw.append([np.inf,  np.inf])
#     w0.append([0, 0])
#     x_plot.append(Xk)

#     # Add equality constraint
#     g.append(Xk_end-Xk)
#     lbg.append([0, 0])
#     ubg.append([0, 0])

# # Concatenate vectors
# w = ca.vertcat(*w)
# g = ca.vertcat(*g)
# x_plot = ca.horzcat(*x_plot)
# u_plot = ca.horzcat(*u_plot)
# w0 = np.concatenate(w0)
# lbw = np.concatenate(lbw)
# ubw = np.concatenate(ubw)
# lbg = np.concatenate(lbg)
# ubg = np.concatenate(ubg)

# # Create an NLP solver
# prob = {'f': J, 'x': w, 'g': g}
# solver = ca.nlpsol('solver', 'ipopt', prob);

# # Function to get x and u trajectories from w
# trajectories = ca.Function('trajectories', [w], [x_plot, u_plot], ['w'], ['x', 'u'])

# # Solve the NLP
# sol = solver(x0=w0, lbx=lbw, ubx=ubw, lbg=lbg, ubg=ubg)
# x_opt, u_opt = trajectories(sol['x'])
# x_opt = x_opt.full() # to numpy array
# u_opt = u_opt.full() # to numpy array

# # Plot the result
# tgrid = np.linspace(0, T, N+1)
# plt.figure(1)
# plt.clf()
# plt.plot(tgrid, x_opt[0], '--')
# plt.plot(tgrid, x_opt[1], '-')
# plt.step(tgrid, np.append(np.nan, u_opt[0]))
# plt.xlabel('t')
# plt.legend(['x1','x2','u'])
# plt.grid()
# plt.show()